"""
文档召回重排
"""
from typing import List
from core.document import Document
from FlagEmbedding import FlagReranker
from threading import RLock
from core.config import settings


class ReRanker(object):

    lock = RLock()

    def __new__(cls, *args, **kwargs):
        with ReRanker.lock:
            if not hasattr(ReRanker, "_instance"):
                ReRanker._instance = object.__new__(cls)
                cls.reranker = FlagReranker(settings.RERANKER_MODEL, use_fp16=True)
        return ReRanker._instance

    def rerank(self, query, documents: List[Document], top_k:int = 10) -> List[Document]:
        sentence_pairs = []
        for doc in documents:
            sentence_pairs.append((query, doc.content))

        scores = self.reranker.compute_score(sentence_pairs, batch_size=4)
        sorted_docs = sorted(zip(scores, documents), key=lambda x: -x[0])
        top_documents = [doc for (score, doc) in sorted_docs[0:top_k]]
        return top_documents